package cn.legym.garp.gateway.traffic;

import cn.legym.garp.gateway.traffic.config.TrafficConfig;
import cn.legym.garp.gateway.traffic.dynamic.DynamicResourceManager;
import cn.legym.garp.gateway.traffic.dynamic.bean.DynamicResourceDefinition;
import cn.legym.garp.gateway.traffic.rule.TrafficRule;
import cn.legym.garp.gateway.traffic.rule.TrafficRuleManager;
import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.EntryType;
import com.alibaba.csp.sentinel.ResourceTypeConstants;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.adapter.gateway.common.SentinelGatewayConstants;
import com.alibaba.csp.sentinel.adapter.gateway.common.param.GatewayParamParser;
import com.alibaba.csp.sentinel.adapter.gateway.sc.ServerWebExchangeItemParser;
import com.alibaba.csp.sentinel.adapter.gateway.sc.exception.SentinelGatewayBlockExceptionHandler;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Set;


/**
 * create by pipisun on 2021/9/6
 * copyright © 2017-2021 Legym Technology Co.,Ltd. All rights reserve
 * d.
 * file description:
 * <p>
 * 自定义动态流控入口
 * 检查请求是否匹配{@link TrafficRule},
 * 通过匹配到的规则触发{@link DynamicResourceDefinition}定义的资源调用
 * <p>
 * last update by {} on {}
 * update description:
 *
 * @see com.alibaba.csp.sentinel.adapter.servlet.CommonFilter
 * @see com.alibaba.csp.sentinel.adapter.gateway.sc.SentinelGatewayFilter
 */
@Component
@Slf4j
public class CustomTrafficControlManager {

    @Autowired
    private TrafficRuleManager trafficRuleManager;
    @Autowired
    private DynamicResourceManager dynamicResourceManager;
    @Autowired
    private TrafficConfig config;
    @Autowired
    private SentinelGatewayBlockExceptionHandler exceptionHandler;

    private final GatewayParamParser<ServerWebExchange> paramParser = new GatewayParamParser<>(
            new ServerWebExchangeItemParser());

    public Mono<Void> filter(ServerWebExchange exchange) {
        if (config == null || !config.isEnabled()) {
            return null;
        }
        try {
            // 匹配请求路径 找到请求转义规则
            Set<TrafficRule> matchedRules = trafficRuleManager.matchRule(exchange.getRequest());
            for (TrafficRule rule : matchedRules) {
                Mono<Void> blocked = triggerResourceIfPresent(exchange, rule);
                if (blocked != null) {
                    return blocked;
                }
            }
//            return matchedRules.stream().map(rule -> triggerResourceIfPresent(exchange, rule));
        } catch (Exception e) {
            log.error("", e);
        }
        return null;
    }

    /**
     * 根据自定义规则匹配结果查找自定义Sentinel资源，并触发限流检查
     *
     * @param exchange    请求
     * @param matchedRule 匹配到的规则
     */
    private Mono<Void> triggerResourceIfPresent(ServerWebExchange exchange, TrafficRule matchedRule) {
        // 转义结果查找资源
        if (matchedRule.getDynamicResource()) { // 自定义动态规则
            boolean resource = dynamicResourceManager.locateResource(exchange.getRequest(), matchedRule);
            if (!resource) {
                log.warn("loading dynamic resource failed. rule=" + matchedRule + ",request=" + exchange.getRequest().getPath().value());
                return null;
            }
        }
        Entry entry = null;
        try {
            Object[] params = paramParser.parseParameterFor(matchedRule.getResource(), exchange,
                    r -> r.getResourceMode() == SentinelGatewayConstants.RESOURCE_MODE_ROUTE_ID);
            // 触发资源访问
            entry = SphU.entry(matchedRule.getResource(), ResourceTypeConstants.COMMON_WEB, EntryType.IN, params);
        } catch (BlockException e) {
            // 访问被限流
            // Return the block page, or redirect to another URL.
//            GatewayCallbackManager.getBlockHandler().handleRequest(exchange, e);
            return exceptionHandler.handle(exchange, e);
        } finally {
            if (entry != null) {
                entry.exit();
            }
        }
        return null;
    }
}
